# -*- coding: utf-8 -*-  
'''
crflinker网络

@author: luoyi
Created on 2021年7月4日
'''
import tensorflow as tf

import utils.conf as conf
import utils.dictionaries as dicts
import utils.relationships as rel
import utils.poses as poses
from models.abstract_nets import AbstractModel
from models.crflinker.losses import CRFLinkerLosser
from models.crflinker.metrics import CRFMetrics, TPLinkerMetrics
from models.bert.part import BertLayer
from models.crflinker.part import CRFLayer, TPLinkerLayer


#    CRF + TPLinker网络
#    用crf层代替tplinker的et握手层
class CRFLinker(AbstractModel):
    def __init__(self,
                 name='crf_linker',
                 #    bert相关参数
                 vocab_size=dicts.dict_size(),
                 n_block=conf.BERT.get_n_block(),
                 n_head=conf.BERT.get_n_head_attention(),
                 d_model=conf.BERT.get_d_model(),
                 f_model=conf.BERT.get_f_model(),
                 dropout_rate=conf.BERT.get_dropout_rate(),
                 #    crflinker相关参数
                 max_sen_len=conf.TPLINKER.get_max_sentence_len(),
                 rel_size=len(rel.id_rel),
                 pos_size=poses.bio_tag_size(),
                 
                 #    训练相关参数
                 batch_size=conf.DATASET_BAIDU.get_batch_size(),
                 loss_lamda_ner=conf.TPLINKER.get_loss_lamda_ner(), 
                 loss_lamda_re=conf.TPLINKER.get_loss_lamda_re(),
                 learning_rate=0.001,
                 input_shape=(None, 2, conf.CRFLINKER.get_max_sentence_len()),
                 
                 auto_assembling=True, 
                 is_build=True,):
        #    bert相关参数
        self._vocab_size = vocab_size
        self._n_block = n_block
        self._n_head = n_head
        self._d_model = d_model
        self._f_model = f_model
        self._dropout_rate = dropout_rate
        #    crflinker相关参数
        self._max_sen_len = max_sen_len
        self._rel_size = rel_size
        self._pos_size = pos_size
        #    模型相关参数
        self._batch_size = batch_size
        self._loss_lamda_ner = loss_lamda_ner
        self._loss_lamda_re = loss_lamda_re
        self._learning_rate = learning_rate
        self._input_shape = input_shape
        
        super(CRFLinker, self).__init__(learning_rate, name, auto_assembling, is_build, input_shape)
        pass
    
    def create_optimizer(self, learning_rate=0.001):
        return tf.optimizers.Adam(learning_rate=learning_rate)
    
    def create_loss(self):
        return CRFLinkerLosser(crf_layer=self._crf_layer, 
                               tplinker_layer=self._tplinker_layer, 
                               loss_lamda_ner=self._loss_lamda_ner, 
                               loss_lamda_re=self._loss_lamda_re)
    
    def create_metrics(self):
        return [CRFMetrics(name='crf_metric', crf_layer=self._crf_layer),
                TPLinkerMetrics(name='tplinker_metric', tplinker_layer=self._tplinker_layer)]
    
    #    装配网络（这步操作其实可以放到build）
    def assembling(self):
        #    bert层
        self._bert_layer = BertLayer(name='bert_layer', 
                                     vocab_size=self._vocab_size, 
                                     max_sen_len=self._max_sen_len, 
                                     max_sen=1+1, 
                                     n_block=self._n_block, 
                                     n_head=self._n_head, 
                                     d_model=self._d_model, 
                                     f_model=self._f_model, 
                                     dropout_rate=self._dropout_rate, 
                                     is_nsp=False, 
                                     is_mlm=False)
        
        #    crf层
        self._crf_layer = CRFLayer(name='crf_layer', 
                                   pos_size=self._pos_size, 
                                   max_sen_len=self._max_sen_len)
        
        #    tplinker层
        self._tplinker_layer = TPLinkerLayer(name='tplinker_layer', 
                                             rel_size=self._rel_size, 
                                             max_sen_len=self._max_sen_len,
                                             d_model=self._d_model)
        pass

    #    前向
    def call(self, inputs, training=None, mask=None):
        #    过bert层
        h = self._bert_layer(inputs)                    #    Tensor(batch_size, max_sen_len, d_model)
        
        #    过crf层
        _ = self._crf_layer(h)                          #    Tensor(batch_size, max_sen_len, pos_size)
        
        #    过tplinker层
        _ = self._tplinker_layer(h)                     #    Tensor(batch_size, rel_size, max_vec_sen, 2, 3)
        
        return h
    
    
    pass
